import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

# Path to the dataset
data_path = '../data/raw_data/334份 按选项序号 汇总变量后.xlsx'

# Function for one-hot encoding if needed
def one_hot_encode(df):
    print('------------独热前-------------')
    print(df.head())
    print('-------------------------')
    column_indices = [0] 
    categorical_columns = df.iloc[:, column_indices].columns.tolist()

    df_encoded = pd.get_dummies(df, columns=categorical_columns)
    print('------------独热后-------------')
    print(df_encoded.head())
    print('-------------------------')

    return df_encoded

# Main analysis
if __name__ == '__main__':
    df = pd.read_excel(data_path)
    
    # Clean data by removing first column
    df = df.drop(df.columns[0], axis=1)

    # 0: 性别 (Gender, categorical data)
    gender = df.iloc[:, 0]
    gender = gender.replace({1: 'Male', 2: 'Female'})

    # Calculate gender distribution percentages
    gender_counts = gender.value_counts(normalize=True)

    # Plot pie chart for gender distribution and save as image
    plt.figure(figsize=(8, 8))
    plt.pie(gender_counts, labels=gender_counts.index, autopct='%1.1f%%', startangle=140)
    plt.title('Gender Distribution')
    plt.savefig('gender_distribution.png')
    plt.close()

    # 1: 年龄 (Age, numerical data)
    age = df.iloc[:, 1]
    
    # Calculate mean, median, and standard deviation
    age_mean = age.mean()
    age_median = age.median()
    age_std = age.std()

    # Plot histogram for age distribution and save as image
    plt.figure(figsize=(8, 6))
    sns.histplot(age, kde=True, bins=20)
    plt.title(f'Age Distribution (Mean: {age_mean:.2f}, Median: {age_median:.2f}, Std: {age_std:.2f})')
    plt.xlabel('Age')
    plt.ylabel('Frequency')
    plt.savefig('age_distribution.png')
    plt.close()

    # 2: 户籍 (Hukou, categorical data)
    hukou = df.iloc[:, 2]
    hukou = hukou.replace({1: 'Urban', 2: 'Rural'})

    # Calculate hukou distribution percentages
    hukou_counts = hukou.value_counts(normalize=True)

    # Plot pie chart for hukou distribution and save as image
    plt.figure(figsize=(8, 8))
    plt.pie(hukou_counts, labels=hukou_counts.index, autopct='%1.1f%%', startangle=140)
    plt.title('Hukou Distribution')
    plt.savefig('hukou_distribution.png')
    plt.close()

    # Output statistical summaries for Age
    print(f"Age - Mean: {age_mean:.2f}, Median: {age_median:.2f}, Std: {age_std:.2f}")
    print(f"Gender Distribution:\n{gender_counts}")
    print(f"Hukou Distribution:\n{hukou_counts}")
